# coding: utf-8

# Adopted from:
#     https://github.com/seshuad/IMagenet/blob/master/TinyImagenet.py

import os
import matplotlib
import numpy as np
import pandas as pd
import tensorflow as tf
import matplotlib.image as mpimg
import cv2
from sklearn import preprocessing
from Params import*


# 获取字典的函数 但是有点怪 参考的版本中用到了一个key
def get_directories():

    IMAGE_DIRECTORY = '/input_dir/datasets/Tiny-ImageNet/'
    TRAINING_IMAGES_DIR = '/input_dir/datasets/Tiny-ImageNet/train/'
    VAL_IMAGES_DIR = '/input_dir/datasets/Tiny-ImageNet/val/'

    return IMAGE_DIRECTORY, TRAINING_IMAGES_DIR, VAL_IMAGES_DIR


def load_training_images(image_dir, batch_size=200):

    image_index = 0

    images = np.ndarray(shape=(NUM_IMAGES, IMAGE_ARR_SIZE))
    names = []
    labels = []
    # print("Loading training images from ", image_dir)

    # Loop through all the types directories
    for type in os.listdir(image_dir):
        if os.path.isdir(image_dir + type + '/images/'):
            type_images = os.listdir(image_dir + type + '/images/')
            # Loop through all the images of a type directory
            batch_index = 0
            # print("Loading Class ", type)
            for image in type_images:
                image_file = os.path.join(image_dir, type + '/images/', image)

                image_data = mpimg.imread(image_file)
                # print('Loaded Image', image_file, image_data.shape)
                '''
                这里遇到了一个问题 有的图片是灰度图而非RGB的三通道彩图
                暂时的处理办法是用opencv的函数将灰度图转为三通道 但不知道是否有其他影响
                '''
                if image_data.shape == (IMAGE_SIZE, IMAGE_SIZE):
                    image_data = cv2.cvtColor(cv2.resize(image_data, (IMAGE_SIZE, IMAGE_SIZE)), cv2.COLOR_GRAY2BGR)

                if image_data.shape == (IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS):
                    images[image_index, :] = image_data.flatten()

                    labels.append(type)
                    names.append(image)

                    image_index += 1
                    batch_index += 1

                if batch_index >= batch_size:
                   break

    print("Loaded Training Images", image_index)
    return images, np.asarray(labels), np.asarray(names)


def get_label_from_name(data, name):
    for idx, row in data.iterrows():
        if row['File'] == name:
            return row['Class']

    return None


def load_validation_images(testdir, validation_data, batch_size=NUM_VAL_IMAGES):
    labels = []
    names = []
    image_index = 0

    images = np.ndarray(shape=(batch_size, IMAGE_ARR_SIZE))
    val_images = os.listdir(testdir + '/images/')

    # Loop through all the images of a val directory
    batch_index = 0

    for image in val_images:
        image_file = os.path.join(testdir, 'images/', image)
        # print (testdir, image_file)

        image_data = mpimg.imread(image_file)
        if image_data.shape == (IMAGE_SIZE, IMAGE_SIZE):
            image_data = cv2.cvtColor(cv2.resize(image_data, (IMAGE_SIZE, IMAGE_SIZE)), cv2.COLOR_GRAY2BGR)

        if image_data.shape == (IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS):
            images[image_index, :] = image_data.flatten()
            image_index += 1
            labels.append(get_label_from_name(validation_data, image))
            names.append(image)
            batch_index += 1
            # print(batch_index)

        if batch_index >= batch_size:
            break

    print("Loaded Validation images ", image_index)
    return images, np.asarray(labels), np.asarray(names)


def reset_graph(seed=42):
    tf.reset_default_graph()
    tf.set_random_seed(seed)
    np.random.seed(seed)